package com.jboltai.scheduler;

import cn.hutool.core.util.StrUtil;
import com.jboltai.capability.embedding.EmbeddingProcessor;
import com.jboltai.config.JBoltAIConfig;
import com.jboltai.event.EventAttrKey;
import com.jboltai.event.EventErrorType;
import com.jboltai.event.EventState;
import com.jboltai.event.embedding.EmbeddingEvent;
import com.jboltai.resource.embedding.EmbeddingResource;
import com.jboltai.resource.embedding.EmbeddingResourceCenter;
import com.jboltai.resource.embedding.EmbeddingResourceGroup;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.Iterator;
import java.util.concurrent.LinkedBlockingQueue;

/**
 * AI事件调度器
 */
class EmbeddingEventScheduler extends Scheduler<EmbeddingEvent>{

    private static final Logger LOGGER = LoggerFactory.getLogger(JBoltAIConfig.LOGGER_NAME);


    /**
     * 每个大模型的事件队列
     */
    private  final LinkedBlockingQueue<EmbeddingEvent> eventQueue = new LinkedBlockingQueue(JBoltAIConfig.getMaxEmbeddingEventWaitCount());

    public EmbeddingEventScheduler() {
        super("embedding");
    }
    /**
     * 发布一个事件
     * @param event
     */
    public void publish(EmbeddingEvent event) {
        event.setState(EventState.WAITING);
        if (event.getEmbeddingModel() != null) {
            boolean hasResource = EmbeddingResourceCenter.checkHasResource(event.getEmbeddingModel(), event.getEmbeddingResourceId());
            if (!hasResource) {
                //没有可用的资源
                if (event.getEmbeddingResourceId() == null) {
                    event.fail(EventErrorType.NO_RESOURCE, "未找到可用的" + event.getEmbeddingModel().getName() + "资源");
                    LOGGER.error("未找到可用的{}资源", event.getEmbeddingModel().getName());
                } else {
                    event.fail(EventErrorType.NO_RESOURCE, "未找到id为" +event.getEmbeddingResourceId()+"的" + event.getEmbeddingModel().getName() + "资源");
                    LOGGER.error("未找到id为{}的{}资源", event.getEmbeddingResourceId(),event.getEmbeddingModel().getName());
                }
                return;
            }
        } else if (StrUtil.isNotBlank(event.getEmbeddingResourceId())) {
            if (EmbeddingResourceCenter.getResourceById(event.getEmbeddingResourceId()) ==  null) {
                event.fail(EventErrorType.NO_RESOURCE, "未找到id为" +event.getEmbeddingResourceId()+"的资源");
                LOGGER.error("未找到id为{}的资源", event.getEmbeddingResourceId());
                return ;
            }

        }


        boolean insert = eventQueue.offer(event);

        if (!insert) {
            //队列满了
            event.fail(EventErrorType.QUEUE_FULL);
        }
    }

    @Override
    void process() {


        try {
            //检测该模型的事件队列 是不是空的
//                if (eventQueue.isEmpty()) {
//                    sleep(20);
//                    continue;
//                }

            Iterator<EmbeddingEvent> iterator = eventQueue.iterator();
            while (iterator.hasNext()) {
                EmbeddingEvent event = iterator.next();
                if (event.getState() == EventState.CANCEL) {
                    //事件被取消了
                    iterator.remove();
                    continue;
                }

                dispatch(event);

                if (event.getState() != EventState.WAITING) {
                    //不是等待状态，就不用再处理了
                    iterator.remove();
                }
            }
//            sleep(10);
        } catch (Exception e) {
            throw new RuntimeException(e);
        }



    }

    @Override
    void dispatch(EmbeddingEvent event) {
        EmbeddingResourceGroup resources = null;
        EmbeddingResource resource = null;

        //embedding事件必须指定模型或者资源id

        if (StrUtil.isNotBlank(event.getEmbeddingResourceId())) {
            //事件指定了某一个资源
            resource = EmbeddingResourceCenter.getResourceById(event.getEmbeddingResourceId());
            if (resource ==  null) {
                event.fail(EventErrorType.NO_RESOURCE, "未找到id为" +event.getEmbeddingResourceId()+"的资源");
                LOGGER.error("未找到id为{}的资源", event.getEmbeddingResourceId());
                return ;
            }

        }
        if (resource == null) {
            //事件指定了 大模型, 根据大模型找资源
            resources = EmbeddingResourceCenter.getModelResources(event.getEmbeddingModel());
            if (resources == null) {
                //没有可用的资源
                event.fail(EventErrorType.NO_RESOURCE, "未找到可用的" + event.getEmbeddingModel().getName() + "资源");
                LOGGER.error("未找到可用的{}资源", event.getEmbeddingModel().getName());
                return;
            }
            resource = resources.getResource(event.getEmbeddingResourceId());
        }

        if (event.getEmbeddingModel() == null && resource != null) {
            //本身没设置用哪个大模型，但是资源有了，就根据资源指定用的大模型
            event.setEmbeddingModel(resource.getModel());
        }

        int remainTimes = 0;
        if (resource == null || (remainTimes= resource.checkAvailableAndReduce()) == -1) {
            //当前没有可用的资源
            if (event.getWaitTimestamp() > 0 && System.currentTimeMillis() > event.getWaitTimestamp()) {
                //超时了还没资源可用，就算了
                event.fail(EventErrorType.QUEUE_WAIT_TIMEOUT);
                LOGGER.error("事件{} 等待超时，放弃执行", event.getId());
            }
            if (event.getState() != EventState.WAITING) {
                //不是等待状态，说明已经处理过一次了，肯定不在等待队列中了
                if (event.getEmbeddingResourceId() == null) {
                    event.fail(EventErrorType.NO_RESOURCE, "未找到可用的" + event.getEmbeddingModel().getName() + "资源");
                    LOGGER.error("当前没有可用的{}资源", event.getEmbeddingModel());
                } else {
                    event.fail(EventErrorType.NO_RESOURCE, "未找到id为" +event.getEmbeddingResourceId()+"的" + event.getEmbeddingModel().getName() + "资源");
                    LOGGER.error("当前没有可用的id为{}的{}资源", event.getEmbeddingResourceId(),event.getEmbeddingModel().getName());
                }
                return;
            }
            return;
        }
        if (remainTimes == 0) {
            //当前资源已经用完了，需要重置
            resources.update(resource);
        }
        if (event.setState(EventState.PROCESSING)) {
            if (event.getText().length > resource.getModel().getChunkCountLimit()) {
                event.fail(EventErrorType.EMBEDDING_EXCEED_CHUNK_COUNT_LIMIT);
                LOGGER.error("当前使用的Embedding大模型，每次可传入的文本片段数不能超过{}，当前传入了{}个", event.getEmbeddingModel().getChunkCountLimit(), event.getText().length);
                return;
            }
            EmbeddingProcessor processor = resource.getModel().getProcessor();
            event.setAttr(EventAttrKey.EVENT_PROCESSOR, processor);
            event.setAttr(EventAttrKey.EVENT_RESOURCE, resource);
            processor.process(event, resource);
        }

    }


}
